# coding=UTF-8

from PIL import Image
from pylab import *
from scipy import ndimage
from scipy.misc import imsave
import pickle
import os



#依次打开目录所有图像，标定坐标并保存到文件
def read_points_from_pic(dir):
    faces={}
    for item in os.listdir(dir):
        if item[-3:]!='jpg':
            continue
        im=array(Image.open(dir+'/'+item))    #载入图像
        imshow(im)
        print('please click 3 points')
        x=ginput(3)
        print('you clicked:')
        faces[item]=hstack(x)
        close()
    return faces;


def compute_rigid_transform(refpoints,points):
    """计算points对应到refpoints的旋转、缩放及平移"""
    A=array([
             [points[0],-points[1],1,0],
             [points[1],points[0],0,1],
             [points[2],-points[3],1,0],
             [points[3],points[3],0,1],
             [points[4],-points[5],1,0],
             [points[5],points[4],0,1]
            ])
    y=array([ refpoints[0],
              refpoints[1],
              refpoints[2],
              refpoints[3],
              refpoints[4],
              refpoints[5] ])
    #y=Ax,  x=[a,b,tx,ty]'
    #最小二乘求解Ax=y，返回的tuple的第一个元素为最小二乘解
    a,b,tx,ty=linalg.lstsq(A,y)[0]
    R=array([[a,-b],[b,a]])
    R=inv(R.T)
    return R,tx,ty

def rigid_alignment(faces,path,plotflag=False):
    """对齐图像，并保存到新图像 """
    refpoints=list(faces.values())[0]   #将第一幅图像中的点做为基准

    for face in faces:  #取字典的key
        points=faces[face]
        #计算仿射变换矩阵
        R,tx,ty=compute_rigid_transform(refpoints,points)
        #T=array([[R[1][1],R[1][0]],[R[0][1],R[0][0]]])
        im=array(Image.open(os.path.join(path,face)))
        im2=zeros(im.shape,'uint8')
        ##进行仿射反变换
        for i in arange(len(im.shape)):
            #Tinv=linalg.inv(T)
            t=matrix(transpose(R))*[[tx],[ty]]
            im2[:, :, i] =ndimage.affine_transform(im[:, :, i],R,
                                                   offset=[-ty,-tx])
        if plotflag:
            imshow(im2)
            show()
        h,w=im2.shape[:2]
        border=(w+h)/20
        ##保存配准后的图像
        imsave(os.path.join(path,'aligned/'+face),
               im2[border:h-border,border:w-border,:])





##############test####################

###########################
# 从系列图像提取点并保存到文件,第一次运行时使用，之后注释掉
#faces=read_points_from_pic('jkfaces')
#file=open("jkfaces/jkfaces.dat",'wb')
#pickle.dump(faces,file)
#file.close()
###########################

file=open("jkfaces/jkfaces.dat",'rb')
faces=pickle.load(file)
file.close()

rigid_alignment(faces,'jkfaces')